# encoding: utf-8
#!/usr/bin/python3
import os
from torch.utils.data import Dataset,DataLoader
import numpy as np
import torch
import json
import cv2
from data_aug import aug_transform,random_concat_image
import time


class MyDataset(Dataset):
    def __init__(self,image_path,json_dir,data_aug = False):
        self.json_dir = json_dir
        self.name_list = json.load(open(json_dir))[:]
        self.image_path = image_path
        self.data_aug = data_aug

    def __getitem__(self, index):
        name = self.name_list[index]
        image_dir = os.path.join(self.image_path,name)
        label_dir = os.path.join(self.image_path,name.replace("tif", "png"))
        image = cv2.imread(image_dir,-1)
        # label为1-10，转成0-9
        label = (cv2.imread(label_dir))[:,:,0] - 1
        if self.data_aug:
            if np.random.uniform() < 0.3:
                start_time = time.time()
                image,label = random_concat_image(self.name_list,self.image_path)
                # print('use time in concat image --',time.time() - start_time)

            aug_data = aug_transform(image=image,mask=label)
            image,label = aug_data['image'],aug_data['mask']
        image = np.transpose(image,(2,0,1))
        # print('image shape', image.shape, 'label shape', label.shape)
        return image,label

    def __len__(self):
        return len(self.name_list)


if __name__ == '__main__':
    dataset = MyDataset(image_path='../suichang_round1_train_210120\suichang_round1_train_210120',json_dir='label_file/train.json',data_aug=True)
    data_loader = DataLoader(dataset=dataset,batch_size=4,num_workers=6)
    for image,label in data_loader:
        image = image.cuda()
        label = label.cuda()
        print(image.shape,label.shape)